#ifndef YOLODETECTOR_H
#define YOLODETECTOR_H

#include <QObject>
#include <opencv2/opencv.hpp>
#include "inc/rknn_api.h"

class YoloDetector : public QObject
{
    Q_OBJECT
public:
    explicit YoloDetector(QObject *parent = nullptr);
    ~YoloDetector();

    bool loadModel(const QString &modelPath);
    bool initRuntime();
    std::vector<cv::Rect> detect(const cv::Mat &image);

private:
    rknn_input_output_num ioNum;
    rknn_context ctx;
    int imgWidth;
    int imgHeight;
    int imgChannel;
    int modelInputWidth;
    int modelInputHeight;
    float boxConfThreshold;
    float nmsThreshold;
    unsigned char *modelData;
    bool modelLoaded;

    bool preprocess(const cv::Mat &image, cv::Mat &resizedImage);
    bool postprocess(const std::vector<float> &output, std::vector<cv::Rect> &boxes);
    void dump_tensor_attr(const rknn_tensor_attr *attr) const;
};

void testYoloDetector();

#endif // YOLODETECTOR_H
